{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "toc": true
   },
   "source": [
    "<h1>Table of Contents<span class=\"tocSkip\"></span></h1>\n",
    "<div class=\"toc\" style=\"margin-top: 1em;\"><ul class=\"toc-item\"></ul></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST Client Example <a class=\"tocSkip\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From : https://github.com/tensorflow/serving/blob/master/tensorflow_serving/example/mnist_input_data.py\n",
    "\"\"\"Functions for downloading and reading MNIST data.\"\"\"\n",
    "\n",
    "from __future__ import print_function\n",
    "\n",
    "import gzip\n",
    "import os\n",
    "\n",
    "import numpy\n",
    "from six.moves import urllib\n",
    "\n",
    "# CVDF mirror of http://yann.lecun.com/exdb/mnist/\n",
    "SOURCE_URL = 'https://storage.googleapis.com/cvdf-datasets/mnist/'\n",
    "TRAIN_IMAGES = 'train-images-idx3-ubyte.gz'\n",
    "TRAIN_LABELS = 'train-labels-idx1-ubyte.gz'\n",
    "TEST_IMAGES = 't10k-images-idx3-ubyte.gz'\n",
    "TEST_LABELS = 't10k-labels-idx1-ubyte.gz'\n",
    "VALIDATION_SIZE = 5000\n",
    "\n",
    "\n",
    "def maybe_download(filename, work_directory):\n",
    "  \"\"\"Download the data from Yann's website, unless it's already here.\"\"\"\n",
    "  if not os.path.exists(work_directory):\n",
    "    os.mkdir(work_directory)\n",
    "  filepath = os.path.join(work_directory, filename)\n",
    "  if not os.path.exists(filepath):\n",
    "    filepath, _ = urllib.request.urlretrieve(SOURCE_URL + filename, filepath)\n",
    "    statinfo = os.stat(filepath)\n",
    "    print('Successfully downloaded %s %d bytes.' % (filename, statinfo.st_size))\n",
    "  return filepath\n",
    "\n",
    "\n",
    "def _read32(bytestream):\n",
    "  dt = numpy.dtype(numpy.uint32).newbyteorder('>')\n",
    "  return numpy.frombuffer(bytestream.read(4), dtype=dt)[0]\n",
    "\n",
    "\n",
    "def extract_images(filename):\n",
    "  \"\"\"Extract the images into a 4D uint8 numpy array [index, y, x, depth].\"\"\"\n",
    "  print('Extracting %s' % filename)\n",
    "  with gzip.open(filename) as bytestream:\n",
    "    magic = _read32(bytestream)\n",
    "    if magic != 2051:\n",
    "      raise ValueError(\n",
    "          'Invalid magic number %d in MNIST image file: %s' %\n",
    "          (magic, filename))\n",
    "    num_images = _read32(bytestream)\n",
    "    rows = _read32(bytestream)\n",
    "    cols = _read32(bytestream)\n",
    "    buf = bytestream.read(rows * cols * num_images)\n",
    "    data = numpy.frombuffer(buf, dtype=numpy.uint8)\n",
    "    data = data.reshape(num_images, rows, cols, 1)\n",
    "    return data\n",
    "\n",
    "\n",
    "def dense_to_one_hot(labels_dense, num_classes=10):\n",
    "  \"\"\"Convert class labels from scalars to one-hot vectors.\"\"\"\n",
    "  num_labels = labels_dense.shape[0]\n",
    "  index_offset = numpy.arange(num_labels) * num_classes\n",
    "  labels_one_hot = numpy.zeros((num_labels, num_classes))\n",
    "  labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1\n",
    "  return labels_one_hot\n",
    "\n",
    "\n",
    "def extract_labels(filename, one_hot=False):\n",
    "  \"\"\"Extract the labels into a 1D uint8 numpy array [index].\"\"\"\n",
    "  print('Extracting %s' % filename)\n",
    "  with gzip.open(filename) as bytestream:\n",
    "    magic = _read32(bytestream)\n",
    "    if magic != 2049:\n",
    "      raise ValueError(\n",
    "          'Invalid magic number %d in MNIST label file: %s' %\n",
    "          (magic, filename))\n",
    "    num_items = _read32(bytestream)\n",
    "    buf = bytestream.read(num_items)\n",
    "    labels = numpy.frombuffer(buf, dtype=numpy.uint8)\n",
    "    if one_hot:\n",
    "      return dense_to_one_hot(labels)\n",
    "    return labels\n",
    "\n",
    "\n",
    "class DataSet(object):\n",
    "  \"\"\"Class encompassing test, validation and training MNIST data set.\"\"\"\n",
    "\n",
    "  def __init__(self, images, labels, fake_data=False, one_hot=False):\n",
    "    \"\"\"Construct a DataSet. one_hot arg is used only if fake_data is true.\"\"\"\n",
    "\n",
    "    if fake_data:\n",
    "      self._num_examples = 10000\n",
    "      self.one_hot = one_hot\n",
    "    else:\n",
    "      assert images.shape[0] == labels.shape[0], (\n",
    "          'images.shape: %s labels.shape: %s' % (images.shape,\n",
    "                                                 labels.shape))\n",
    "      self._num_examples = images.shape[0]\n",
    "\n",
    "      # Convert shape from [num examples, rows, columns, depth]\n",
    "      # to [num examples, rows*columns] (assuming depth == 1)\n",
    "      assert images.shape[3] == 1\n",
    "      images = images.reshape(images.shape[0],\n",
    "                              images.shape[1] * images.shape[2])\n",
    "      # Convert from [0, 255] -> [0.0, 1.0].\n",
    "      images = images.astype(numpy.float32)\n",
    "      images = numpy.multiply(images, 1.0 / 255.0)\n",
    "    self._images = images\n",
    "    self._labels = labels\n",
    "    self._epochs_completed = 0\n",
    "    self._index_in_epoch = 0\n",
    "\n",
    "  @property\n",
    "  def images(self):\n",
    "    return self._images\n",
    "\n",
    "  @property\n",
    "  def labels(self):\n",
    "    return self._labels\n",
    "\n",
    "  @property\n",
    "  def num_examples(self):\n",
    "    return self._num_examples\n",
    "\n",
    "  @property\n",
    "  def epochs_completed(self):\n",
    "    return self._epochs_completed\n",
    "\n",
    "  def next_batch(self, batch_size, fake_data=False):\n",
    "    \"\"\"Return the next `batch_size` examples from this data set.\"\"\"\n",
    "    if fake_data:\n",
    "      fake_image = [1] * 784\n",
    "      if self.one_hot:\n",
    "        fake_label = [1] + [0] * 9\n",
    "      else:\n",
    "        fake_label = 0\n",
    "      return [fake_image for _ in range(batch_size)], [\n",
    "          fake_label for _ in range(batch_size)\n",
    "      ]\n",
    "    start = self._index_in_epoch\n",
    "    self._index_in_epoch += batch_size\n",
    "    if self._index_in_epoch > self._num_examples:\n",
    "      # Finished epoch\n",
    "      self._epochs_completed += 1\n",
    "      # Shuffle the data\n",
    "      perm = numpy.arange(self._num_examples)\n",
    "      numpy.random.shuffle(perm)\n",
    "      self._images = self._images[perm]\n",
    "      self._labels = self._labels[perm]\n",
    "      # Start next epoch\n",
    "      start = 0\n",
    "      self._index_in_epoch = batch_size\n",
    "      assert batch_size <= self._num_examples\n",
    "    end = self._index_in_epoch\n",
    "    return self._images[start:end], self._labels[start:end]\n",
    "\n",
    "\n",
    "def read_data_sets(train_dir, fake_data=False, one_hot=False):\n",
    "  \"\"\"Return training, validation and testing data sets.\"\"\"\n",
    "\n",
    "  class DataSets(object):\n",
    "    pass\n",
    "\n",
    "  data_sets = DataSets()\n",
    "\n",
    "  if fake_data:\n",
    "    data_sets.train = DataSet([], [], fake_data=True, one_hot=one_hot)\n",
    "    data_sets.validation = DataSet([], [], fake_data=True, one_hot=one_hot)\n",
    "    data_sets.test = DataSet([], [], fake_data=True, one_hot=one_hot)\n",
    "    return data_sets\n",
    "\n",
    "  local_file = maybe_download(TRAIN_IMAGES, train_dir)\n",
    "  train_images = extract_images(local_file)\n",
    "\n",
    "  local_file = maybe_download(TRAIN_LABELS, train_dir)\n",
    "  train_labels = extract_labels(local_file, one_hot=one_hot)\n",
    "\n",
    "  local_file = maybe_download(TEST_IMAGES, train_dir)\n",
    "  test_images = extract_images(local_file)\n",
    "\n",
    "  local_file = maybe_download(TEST_LABELS, train_dir)\n",
    "  test_labels = extract_labels(local_file, one_hot=one_hot)\n",
    "\n",
    "  validation_images = train_images[:VALIDATION_SIZE]\n",
    "  validation_labels = train_labels[:VALIDATION_SIZE]\n",
    "  train_images = train_images[VALIDATION_SIZE:]\n",
    "  train_labels = train_labels[VALIDATION_SIZE:]\n",
    "\n",
    "  data_sets.train = DataSet(train_images, train_labels)\n",
    "  data_sets.validation = DataSet(validation_images, validation_labels)\n",
    "  data_sets.test = DataSet(test_images, test_labels)\n",
    "\n",
    "  return data_sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From https://github.com/tensorflow/serving/blob/master/tensorflow_serving/example/mnist_client.py\n",
    "\n",
    "from __future__ import print_function\n",
    "\n",
    "import sys\n",
    "import threading\n",
    "\n",
    "from grpc.beta import implementations\n",
    "import numpy\n",
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow_serving.apis import predict_pb2\n",
    "from tensorflow_serving.apis import prediction_service_pb2\n",
    "\n",
    "class _ResultCounter(object):\n",
    "  \"\"\"Counter for the prediction results.\"\"\"\n",
    "\n",
    "  def __init__(self, num_tests, concurrency):\n",
    "    self._num_tests = num_tests\n",
    "    self._concurrency = concurrency\n",
    "    self._error = 0\n",
    "    self._done = 0\n",
    "    self._active = 0\n",
    "    self._condition = threading.Condition()\n",
    "\n",
    "  def inc_error(self):\n",
    "    with self._condition:\n",
    "      self._error += 1\n",
    "\n",
    "  def inc_done(self):\n",
    "    with self._condition:\n",
    "      self._done += 1\n",
    "      self._condition.notify()\n",
    "\n",
    "  def dec_active(self):\n",
    "    with self._condition:\n",
    "      self._active -= 1\n",
    "      self._condition.notify()\n",
    "\n",
    "  def get_error_rate(self):\n",
    "    with self._condition:\n",
    "      while self._done != self._num_tests:\n",
    "        self._condition.wait()\n",
    "      return self._error / float(self._num_tests)\n",
    "\n",
    "  def throttle(self):\n",
    "    with self._condition:\n",
    "      while self._active == self._concurrency:\n",
    "        self._condition.wait()\n",
    "      self._active += 1\n",
    "\n",
    "\n",
    "def _create_rpc_callback(label, result_counter):\n",
    "  \"\"\"Creates RPC callback function.\n",
    "  Args:\n",
    "    label: The correct label for the predicted example.\n",
    "    result_counter: Counter for the prediction result.\n",
    "  Returns:\n",
    "    The callback function.\n",
    "  \"\"\"\n",
    "  def _callback(result_future):\n",
    "    \"\"\"Callback function.\n",
    "    Calculates the statistics for the prediction result.\n",
    "    Args:\n",
    "      result_future: Result future of the RPC.\n",
    "    \"\"\"\n",
    "    exception = result_future.exception()\n",
    "    if exception:\n",
    "      result_counter.inc_error()\n",
    "      print(exception)\n",
    "    else:\n",
    "      sys.stdout.write('.')\n",
    "      sys.stdout.flush()\n",
    "      response = numpy.array(\n",
    "          result_future.result().outputs['outputs'].float_val)\n",
    "      prediction = numpy.argmax(response)\n",
    "      if label != prediction:\n",
    "        result_counter.inc_error()\n",
    "    result_counter.inc_done()\n",
    "    result_counter.dec_active()\n",
    "  return _callback\n",
    "\n",
    "\n",
    "def do_inference(hostport, work_dir, concurrency, num_tests):\n",
    "  \"\"\"Tests PredictionService with concurrent requests.\n",
    "  Args:\n",
    "    hostport: Host:port address of the PredictionService.\n",
    "    work_dir: The full path of working directory for test data set.\n",
    "    concurrency: Maximum number of concurrent requests.\n",
    "    num_tests: Number of test images to use.\n",
    "  Returns:\n",
    "    The classification error rate.\n",
    "  Raises:\n",
    "    IOError: An error occurred processing test data set.\n",
    "  \"\"\"\n",
    "  test_data_set = read_data_sets(work_dir).test #mnist_input_data.read_data_sets(work_dir).test\n",
    "  host, port = hostport.split(':')\n",
    "  channel = implementations.insecure_channel(host, int(port))\n",
    "  stub = prediction_service_pb2.beta_create_PredictionService_stub(channel)\n",
    "  result_counter = _ResultCounter(num_tests, concurrency)\n",
    "  for i in range(num_tests):\n",
    "    request = predict_pb2.PredictRequest()\n",
    "    request.model_spec.name = 'mnist'\n",
    "    request.model_spec.signature_name = 'predict_images'\n",
    "    image, label = test_data_set.next_batch(1)\n",
    "    request.inputs['inputs'].CopyFrom(\n",
    "        tf.contrib.util.make_tensor_proto(image[0], shape=[1, image[0].size]))\n",
    "    result_counter.throttle()\n",
    "    result_future = stub.Predict.future(request, 5.0)  # 5 seconds\n",
    "    result_future.add_done_callback(\n",
    "        _create_rpc_callback(label[0], result_counter))\n",
    "    if i%50==0:\n",
    "         sys.stdout.write('\\n')\n",
    "  return result_counter.get_error_rate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting /home/armando/datasets/mnist/train-images-idx3-ubyte.gz\n",
      "Extracting /home/armando/datasets/mnist/train-labels-idx1-ubyte.gz\n",
      "Extracting /home/armando/datasets/mnist/t10k-images-idx3-ubyte.gz\n",
      "Extracting /home/armando/datasets/mnist/t10k-labels-idx1-ubyte.gz\n",
      "\n",
      "..................................................\n",
      "..................................................\n",
      "Inference error rate: 7.0%\n"
     ]
    }
   ],
   "source": [
    "error_rate = do_inference(hostport='0.0.0.0:8500', \n",
    "                          work_dir='/home/armando/datasets/mnist',\n",
    "                          concurrency=1, \n",
    "                          num_tests=100)\n",
    "print('\\nInference error rate: %s%%' % (error_rate * 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "toc_cell": true,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
